#!/usr/bin/env python

# --------------------------------------------------------
# Faster R-CNN
# Copyright (c) 2015 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# Written by Ross Girshick
# --------------------------------------------------------

"""
Demo script showing detections in sample images.

See README.md for installation instructions before running.
"""

import _init_paths
from fast_rcnn.config import cfg
from fast_rcnn.test import im_detect
from fast_rcnn.nms_wrapper import nms
from utils.timer import Timer
import numpy as np
import scipy.io as sio
import caffe, os, sys, cv2, copy
import argparse

NETS = {'vgg16': ('VGG16',
                  'VGG16_faster_rcnn_final.caffemodel'),
        'zf': ('ZF',
                  'ZF_faster_rcnn_final.caffemodel')}


def vis_detections(im, class_name, dets, thresh=0.5):
    """Draw detected bounding boxes."""
    #inds = np.where(dets[:, -1] >= thresh)[0]
    #if len(inds) == 0:
    #    return
    w=im.shape[0]
    h=im.shape[1]
    img=copy.deepcopy(im)
    bbox = dets[0, :4]
    score = dets[0, -1]
    x=bbox[0]
    y=bbox[1]
    xx=bbox[2]
    yy=bbox[3]
    cv2.rectangle(img,(x,y),(xx,yy),(0,0,255),2)
    cv2.imshow('Image',img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
        #ax.add_patch(
        #    plt.Rectangle((bbox[0], bbox[1]),
        #                  bbox[2] - bbox[0],
        #                  bbox[3] - bbox[1], fill=False,
        #                  edgecolor='red', linewidth=3.5)
        #    )

       
    return img,x,y,xx,yy

def detect_img(net, image_name):
    im_file = os.path.join(cfg.DATA_DIR, 'demo', image_name)
    im = cv2.imread(im_file)

    # Detect all object classes and regress object bounds
    scores, boxes = im_detect(net, im)
    return scores,boxes,im

def parse_args():
    """Parse input arguments."""
    parser = argparse.ArgumentParser(description='Faster R-CNN demo')
    parser.add_argument('--gpu', dest='gpu_id', help='GPU device id to use [0]',
                        default=0, type=int)
    parser.add_argument('--cpu', dest='cpu_mode',
                        help='Use CPU mode (overrides --gpu)',
                        action='store_true')
    parser.add_argument('--net', dest='demo_net', help='Network to use [vgg16]',
                        choices=NETS.keys(), default='vgg16')

    args = parser.parse_args()

    return args

if __name__ == '__main__':
    cfg.TEST.HAS_RPN = True  # Use RPN for proposals

    args = parse_args()

    prototxt = os.path.join(cfg.MODELS_DIR, NETS[args.demo_net][0],
                            'faster_rcnn_alt_opt', 'faster_rcnn_test.pt')
    caffemodel = os.path.join(cfg.DATA_DIR, 'faster_rcnn_models',
                              NETS[args.demo_net][1])

    if not os.path.isfile(caffemodel):
        raise IOError(('{:s} not found.\nDid you run ./data/script/'
                       'fetch_faster_rcnn_models.sh?').format(caffemodel))


    if args.cpu_mode:
        caffe.set_mode_cpu()
    else:
        caffe.set_mode_gpu()
        caffe.set_device(args.gpu_id)
        cfg.GPU_ID = args.gpu_id
    net = caffe.Net(prototxt, caffemodel, caffe.TEST)

    im_names = ['GT02.png']
    for im_name in im_names:
        print '~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
        print 'Demo for data/demo/{}'.format(im_name)
        scores,boxes,im=detect_img(net, im_name)
    m=boxes.shape[0]
    n=boxes.shape[1]
    group_n=int(n/4)


#set 1 class as example
    group_box=boxes[:,0:4]
    group_score=scores[:,0]
    dets = np.hstack((group_box,group_score[:, np.newaxis])).astype(np.float32)
    #img_width,img_height,x,y,xx,yy=vis_detections(im, 'obj', dets[:1,:], thresh=0.0)
    rec,x,y,xx,yy=vis_detections(im, 'obj', dets[:1,:], thresh=0.0)
    cv2.imshow('Image',im)
    cv2.waitKey(0)
    #print img_width,img_height,x,y,xx,yy
    img=im[int(y):int(yy),int(x):int(xx)]
    cv2.imshow('new_image',img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()